#include "socket.h"

#if defined(AE_WIN32)
#undef errno
#define errno WSAGetLastError()
typedef DWORD __errno_t;
#define SHUT_WR SD_BOTH
#define EWOULDBLOCK WSAEWOULDBLOCK
#define ETIMEDOUT WSAETIMEDOUT
#else
typedef int __errno_t;
static const int one = 1;
#endif

static inline int address_size(const socket_addr *raddr)
{
	if(raddr->sa.sa_family == AF_INET)
		return sizeof(struct sockaddr_in);
	else if(raddr->sa.sa_family == AF_INET6)
		return sizeof(struct sockaddr_in6);
	return 0;
}

static inline void watch_event(socket_engine* rsock, int type)
{
	struct poll_customer* customer = TO_POLL_CUSTOMER(rsock);
	customer->__provider->__method->watch_event(customer, type);
}

static int socket_wait_read(socket_engine* rsock)
{
	int retval = -1;
	if(atomic_read(&rsock->nr_closed) == 0) {
#if !defined(AE_WIN32)
		retval = 0;
#else
		if(rsock->flags.is_listen) {
			rsock->adata.pending = socket(rsock->family, SOCK_STREAM, IPPROTO_TCP);
			if(rsock->adata.pending != INVALID_SOCKET) {
				if(AcceptEx(rsock->sock, rsock->adata.pending, rsock->adata.addr_buf, 0, sizeof(rsock->adata.addr_buf) / 2, sizeof(rsock->adata.addr_buf) / 2, &rsock->adata.addr_bytes, (LPWSAOVERLAPPED)&rsock->adata.ola) || WSAGetLastError() == ERROR_IO_PENDING)
					retval = 0;
				else {
					closesocket(rsock->adata.pending);
					rsock->adata.pending = INVALID_SOCKET;
				}
			}
		} else {
			if(rsock->type == SOCK_STREAM) {
				rsock->rwdata.olr.flags = 0;
				if(WSARecv(rsock->sock, &rsock->rwdata.buf, 1, &rsock->rwdata.olr.bytes, &rsock->rwdata.olr.flags, (LPWSAOVERLAPPED)&rsock->rwdata.olr, NULL) != SOCKET_ERROR || WSAGetLastError() == ERROR_IO_PENDING)
					retval = 0;
			}
			else if(rsock->type == SOCK_DGRAM) {
				rsock->rwdata.olr.flags = MSG_PEEK;
				if(WSARecvFrom(rsock->sock, &rsock->rwdata.buf, 1, &rsock->rwdata.olr.bytes, &rsock->rwdata.olr.flags, NULL, 0, (LPWSAOVERLAPPED)&rsock->rwdata.olr, NULL) != SOCKET_ERROR || WSAGetLastError() == ERROR_IO_PENDING)
					retval = 0;
			}
		}
		if(retval == 0)
			atomic_inc(&rsock->nr_managed);
#endif
	}
	if(retval == 0)
		watch_event(rsock, AE_POLL_READ);
	return retval;
}

static int socket_wait_write(socket_engine* rsock, socket_addr* raddr)
{
	int retval = -1;
	if(atomic_read(&rsock->nr_closed) == 0) {
#if !defined(AE_WIN32)
		retval = 0;
#else
		if(rsock->type == SOCK_STREAM) {
			if(WSASend(rsock->sock, &rsock->rwdata.buf, 1, &rsock->rwdata.olw.bytes, 0, (LPWSAOVERLAPPED)&rsock->rwdata.olw, NULL) != SOCKET_ERROR || WSAGetLastError() == ERROR_IO_PENDING)
				retval = 0;
		}
		else if(rsock->type == SOCK_DGRAM) {
			if(WSASendTo(rsock->sock, &rsock->rwdata.buf, 1, &rsock->rwdata.olw.bytes, MSG_DONTROUTE, &raddr->sa, sizeof(*raddr), (LPWSAOVERLAPPED)&rsock->rwdata.olw, NULL) != SOCKET_ERROR || WSAGetLastError() == ERROR_IO_PENDING)
				retval = 0;
		}
		if(retval == 0)
			atomic_inc(&rsock->nr_managed);
#endif
	}
	if(retval == 0)
		watch_event(rsock, AE_POLL_WRITE);
	return retval;
}

static void socket_event_arrived(struct poll_customer* customer, int type)
{
#if defined(AE_WIN32)
	socket_engine * rsock = (socket_engine*) customer;

	if(type == AE_POLL_READ) {
		if(rsock->flags.is_listen) {
			rsock->adata.ready = rsock->adata.pending;
			if(rsock->adata.ready != INVALID_SOCKET) {
				rsock->adata.pending = INVALID_SOCKET;
				setsockopt(rsock->adata.ready, SOL_SOCKET, SO_UPDATE_ACCEPT_CONTEXT, (char *)&rsock->sock, sizeof(rsock->sock));
			}
		}
	}
	atomic_dec(&rsock->nr_managed);
#endif
}

static int socket_try_destory(socket_engine * rsock)
{
#if defined(AE_WIN32) || defined(AE_FREEBSD)
	if(atomic_read(&rsock->nr_managed) != 0)
		return -1;
#endif

#if defined(AE_WIN32)
	if(rsock->flags.is_listen) {
		if(rsock->adata.pending != INVALID_SOCKET)
			closesocket(rsock->adata.pending);
		if(rsock->adata.ready != INVALID_SOCKET)
			closesocket(rsock->adata.ready);
	}
#endif

	closesocket(rsock->sock);
	ae_free(rsock);
	return 0;
}

static int socket_mantain(struct poll_customer* customer, int is_down)
{
	socket_engine * rsock = (socket_engine*) customer;
	if(is_down) {
#if defined(AE_FREEBSD)
		atomic_dec(rsock->nr_managed);
#endif
	}
	if(atomic_read(&rsock->nr_closed) == 0)
		return -1;
	return socket_try_destory(rsock);
}

static inline int set_noblocking(socket_engine * rsock, int on)
{
#if defined(AE_WIN32)
	u_long mode = on;
	if(ioctlsocket(rsock->sock, FIONBIO, &mode) != NO_ERROR)
		return -1;
#else
	if(ioctl(rsock->sock, FIONBIO, (char*)&on) != 0)
		return -1;
#endif
	return 0;
}

static int socket_before_attath(struct poll_customer* customer,__poll_t* handle,int* events)
{
	socket_engine * rsock = (socket_engine*) customer;
	if(set_noblocking(rsock, 1) != 0)
		return -1;
	*handle = (__poll_t)rsock->sock;
	*events = AE_POLL_READ;
	if(!rsock->flags.is_listen) {
		*events |= AE_POLL_WRITE;
		if(rsock->type == SOCK_STREAM)
			*events |= AE_POLL_HANGUP;
	}
#if defined(AE_FREEBSD)
	atomic_set(rsock->nr_managed, 2);
#endif
	return 0;
}

static void socket_after_attath(struct poll_customer* customer,int failed)
{
	socket_engine * rsock = (socket_engine*) customer;
	if(failed) {
		set_noblocking(rsock, 0);
#if defined(AE_FREEBSD)
		atomic_set(rsock->nr_managed, 0);
#endif
	}
	else {
#if defined(AE_WIN32)
		socket_wait_read(rsock);
		socket_wait_write(rsock, NULL);
#endif
	}
}

static struct poll_customer_method socket_method = {
	socket_mantain,
	socket_before_attath,
	socket_after_attath,
	socket_event_arrived
};

static inline void socket_init(socket_engine* rsock)
{
	SET_OBJECT_TYPE(rsock, OBJ_SOCKET);
	TO_POLL_CUSTOMER(rsock)->__method = &socket_method;
#if defined(AE_WIN32)
	if(rsock->flags.is_listen) {
		rsock->adata.pending = INVALID_SOCKET;
		rsock->adata.ready = INVALID_SOCKET;
		rsock->adata.ola.type = AE_POLL_READ;
	}
	else {
		rsock->rwdata.olr.type = AE_POLL_READ;
		rsock->rwdata.olw.type = AE_POLL_WRITE;
	}
#endif
}

int ae_socket_close(ae_object_t sock)
{
	socket_engine* rsock = (socket_engine*)sock;
	struct poll_customer* customer = TO_POLL_CUSTOMER(rsock);

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || atomic_cmpxchg(&rsock->nr_closed, 0, 1) != 0)
		return -1;

	if(customer->__provider)
		customer->__provider->__method->deatch(customer);

	shutdown(rsock->sock, SHUT_WR);
	socket_try_destory(rsock);
	return 0;
}

int ae_socket_async_timeout(ae_object_t sock, int seconds, int type)
{
	struct poll_customer* customer;
	socket_engine* rsock = (socket_engine*)sock;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || seconds < 0)
		return -1;

	customer = TO_POLL_CUSTOMER(rsock);
	if(type == AE_TYPE_IN)
		customer->read_node.seconds = seconds;
	else
		customer->write_node.seconds = seconds;
	if(customer->__provider)
		customer->__provider->__method->reset_timeout(customer, type == AE_TYPE_IN ? AE_POLL_READ: AE_POLL_WRITE);
	return 0;
}

int ae_stream_listen(ae_object_t* sock, const ae_address_t* addr, int backlog)
{
	int retval = -1;
	socket_engine* rsock;
	const socket_addr* raddr = (const socket_addr*)addr;

	if(sock == NULL || raddr == NULL)
		return retval;
	rsock = (socket_engine*)ae_mallocz(sizeof(socket_engine));
	if(rsock == NULL)
		return retval;
	rsock->type = SOCK_STREAM;
	rsock->family = raddr->sa.sa_family;
	rsock->sock = socket(raddr->sa.sa_family, SOCK_STREAM, IPPROTO_TCP);
	if(rsock->sock != INVALID_SOCKET) {
#if !defined(AE_WIN32)
		setsockopt(rsock->sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#endif
		if(bind(rsock->sock, &raddr->sa, address_size(raddr)) == 0
			&& listen(rsock->sock, backlog) == 0){
				rsock->flags.is_listen = 1;
				retval = 0;
		}
		if(retval != 0)
			closesocket(rsock->sock);
	}

	if(retval != 0)
		ae_free(rsock);
	else {
		socket_init(rsock);
		*sock = (ae_object_t)rsock;
	}
	return retval;
}

int ae_stream_accept(ae_object_t sock, ae_object_t* child)
{
	int retval = -1;
	socket_engine* rchild;
	socket_engine* rsock = (socket_engine*)sock;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || !rsock->flags.is_listen)
		return retval;

	rchild = (socket_engine*)ae_mallocz(sizeof(socket_engine));
	if(rchild == NULL)
		return retval;

	rchild->type = SOCK_STREAM;
	rchild->family = rsock->family;
#if defined(AE_WIN32)
	if(rsock->adata.ready != INVALID_SOCKET) {
		rchild->sock = rsock->adata.ready;
		rsock->adata.ready = INVALID_SOCKET;
	}
	else
		rchild->sock = accept(rsock->sock, NULL, NULL);
#else
	rchild->sock = accept(rsock->sock, NULL, NULL);
#endif

	if(rchild->sock == INVALID_SOCKET) {
		__errno_t err = errno;
		if(err == EWOULDBLOCK) {
			if(socket_wait_read(rsock) != 0)
				retval = -2;
		}
		ae_free(rchild);
	}
	else {
		retval = 0;
		socket_init(rchild);
		*child = (ae_object_t)rchild;
	}
	return retval;
}

int ae_stream_connect(ae_object_t* sock, const ae_address_t* addr)
{
	int retval = -1;
	socket_engine* rsock;
	const socket_addr* raddr = (const socket_addr*)addr;

	if(sock == NULL || raddr == NULL)
		return retval;

	rsock = (socket_engine*)ae_mallocz(sizeof(socket_engine));
	if(rsock == NULL)
		return retval;

	rsock->type = SOCK_STREAM;
	rsock->family = raddr->sa.sa_family;
	rsock->sock = socket(raddr->sa.sa_family, SOCK_STREAM, IPPROTO_TCP);
	if(rsock->sock != INVALID_SOCKET) {
		if(connect(rsock->sock, &raddr->sa, address_size(raddr)) == 0)
			retval = 0;
		else
			closesocket(rsock->sock);
	}

	if(retval != 0)
		ae_free(rsock);
	else {
		socket_init(rsock);
		*sock = (ae_object_t)rsock;
	}
	return retval;
}

int ae_stream_recv(ae_object_t sock, void* buffer, int size)
{
	int retval = -1;
	socket_engine* rsock = (socket_engine*)sock;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || buffer == NULL || size <1 || rsock->type != SOCK_STREAM)
		return retval;

	retval = recv(rsock->sock, (char*)buffer, size, 0);
	if(retval < 0) {
		__errno_t err = errno;
		if(err == EWOULDBLOCK) {
			if(socket_wait_read(rsock) != 0)
				retval = 0;
		}
		else if(err != ETIMEDOUT)
			retval = 0;
	}
	return retval;
}

int ae_stream_send(ae_object_t sock, const void* buffer, int size)
{
	int retval = -1;
	socket_engine* rsock = (socket_engine*)sock;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || buffer == NULL || size <1 || rsock->type != SOCK_STREAM)
		return retval;

	retval = send(rsock->sock, (const char*)buffer, size, 0);
	if(retval < 0) {
		__errno_t err = errno;
		if(err == EWOULDBLOCK) {
			if(socket_wait_write(rsock, NULL) != 0)
				retval = 0;
		}
		else if(err != ETIMEDOUT)
			retval = 0;
	}
	return retval;
}

int ae_dgram_bind(ae_object_t* sock, const ae_address_t* addr)
{
	int retval = -1;
	socket_engine* rsock;
	const socket_addr* raddr = (const socket_addr*)addr;

	if(sock == NULL || raddr == NULL)
		return retval;

	rsock = (socket_engine*)ae_mallocz(sizeof(socket_engine));
	if(rsock == NULL)
		return retval;

	rsock->type = SOCK_DGRAM;
	rsock->family = raddr->sa.sa_family;
	rsock->sock = socket(raddr->sa.sa_family, SOCK_DGRAM, IPPROTO_UDP);
	if(rsock->sock != INVALID_SOCKET) {
#if !defined(AE_WIN32)
		setsockopt(rsock->sock, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one));
#endif
		if(bind(rsock->sock, &raddr->sa, address_size(raddr)) == 0)
			retval = 0;
		else
			closesocket(rsock->sock);
	}

	if(retval != 0)
		ae_free(rsock);
	else {
		socket_init(rsock);
		*sock = (ae_object_t)rsock;
	}
	return retval;
}

int ae_dgram_recv(ae_object_t sock, void* buffer, int size, ae_address_t* source)
{
	int retval = -1;
	socklen_t addrlen;
	socket_engine* rsock = (socket_engine*)sock;
	socket_addr* raddr = (socket_addr*)source;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || buffer == NULL || size <1 || raddr == NULL || rsock->type != SOCK_DGRAM)
		return retval;

	addrlen = sizeof(socket_addr);
	retval = recvfrom(rsock->sock, (char*)buffer, size, 0, &raddr->sa, &addrlen);
	if(retval < 0) {
		__errno_t err = errno;
		if(err == EWOULDBLOCK) {
			if(socket_wait_read(rsock) != 0)
				retval = -2;
		}
		else if(err != ETIMEDOUT)
			retval = -2;
	}
	return retval;
}

int ae_dgram_send(ae_object_t sock, const void* buffer, int size, ae_address_t* target)
{
	int retval = -1;
	socket_engine* rsock = (socket_engine*)sock;
	socket_addr* raddr = (socket_addr*)target;

	if(!CHECK_OBJECT(rsock, OBJ_SOCKET) || buffer == NULL || size <1 || raddr == NULL || rsock->type != SOCK_DGRAM)
		return retval;

	retval = sendto(rsock->sock, (const char*)buffer, size, 0, &raddr->sa, address_size(raddr));
	if(retval < 0) {
		__errno_t err = errno;
		if(err == EWOULDBLOCK) {
			if(socket_wait_write(rsock, raddr) != 0)
				retval = -2;
		}
		else if(err != ETIMEDOUT)
			retval = -2;
	}
	return retval;
}
